#include <device_launch_parameters.h>

#include "build_kernels_cuda.h"
#include <Utility/math_utils.h>
#include <Geometry/vector.h>

namespace shermit {
namespace lbvh {
namespace GPU {

	using uint = unsigned int;

	template<class BuildManager_t, class BVSOA_t>
	__global__ void buildBVs(BuildManager_t object_manager, BVSOA_t bvs) {
		int idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= object_manager.primSize()) return;

		//leaves.primIdx[idx] = idx;
		auto bv = object_manager.buildBV(idx);
		bv.setBV(idx, bv);
	}

	template<class BVSOA_t>
	__global__ void calcSceneBV_1(int size, BVSOA_t bv_soa, typename BVSOA_t::bv_struct_t* d_out) {
		uint idx = threadIdx.x + blockIdx.x * blockDim.x;
		uint tid = threadIdx.x;
		extern __shared__ typename BVSOA_t::bv_struct_t sdata[]; //use shared mem to avoid writing global mem
		if (idx < size) sdata[tid] = bv_soa.getBV(idx);
		//if (idx < size) sdata[tid] = leaves.bv[idx];
		else sdata[tid] = typename BVSOA_t::bv_struct_t();
		__syncthreads();

		for (uint s = blockDim.x / 2; s > 0; s >>= 1) {
			if (tid < s) {
				sdata[tid].merge(sdata[tid + s]);
			}
			__syncthreads(); //ensure all merges at one iteration are done
		}
		if (tid == 0) {
			d_out[blockIdx.x] = sdata[tid];
		}
	}

	template<class BVSOA_t>
	__global__ void calcSceneBV_2(int size, BVSOA_t* d_in, typename BVSOA_t::bv_struct_t* d_out) {
		uint idx = threadIdx.x + blockIdx.x * blockDim.x;
		uint tid = threadIdx.x;
		extern __shared__ typename BVSOA_t::bv_struct_t sdata[]; //use shared mem to avoid writing global mem
		if (idx < size) sdata[tid] = d_in[idx];
		else sdata[tid] = typename BVSOA_t::bv_struct_t();
		__syncthreads();

		for (uint s = blockDim.x / 2; s > 0; s >>= 1) {
			if (tid < s) {
				sdata[tid].merge(sdata[tid + s]);
			}
			__syncthreads(); //ensure all merges at one iteration are done
		}
		if (tid == 0) {
			d_out[blockIdx.x] = sdata[tid];
		}
	}

	template<typename Float_t>
	__global__ void calcMortonCodes_32b<AABB2_SOA<Float_t>>(int size, LBvhLeaves<AABB2_SOA<Float_t>> leaves, AABB2<Float_t>* scene) {
		uint idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= size) return;

		const auto offset = (leaves.bv.center(idx) - scene->ptMin);
		offset.x /= (scene->ptMax.x - scene->ptMin.x);
		offset.y /= (scene->ptMax.y - scene->ptMin.y);
		uint mc32 = morton2f_32b(offset.x, offset.y);
		leaves.morton[idx] = mc32;
	}

	template<typename BVSOA_t>
	__global__ void reorderLeaves(int size, LBvhLeaves<BVSOA_t> leaves, LBvhLeaves<BVSOA_t> leaves_sorted, int* mapidx) {
		uint idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= size) return;

		int idx_old = mapidx[idx];
		leaves_sorted.primIdx[idx] = idx_old;
		auto bv = leaves.bv.getBV(idx_old);
		leaves_sorted.bv.setBV(idx, bv);
		//leaves_sorted.bv[idx] = leaves.bv[idx_old];
		leaves_sorted.morton[idx] = leaves.morton[idx];
	}

	template<typename BVSOA_t>
	__global__ void calcBuildMetrics(int interiorSize, LBvhLeaves<BVSOA_t> leaves, uint* metrics) {
		uint idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= interiorSize) return;

		// reflects the highest differing bit between the keys covered by interior node idx
		metrics[idx] = 32 - __clz(leaves.morton[idx] ^ leaves.morton[idx + 1]);
	}

	template<typename BVSOA_t>
	__global__ void buildInteriors(int leafSize, LBvhLeaves<BVSOA_t> leaves, LBvhInteriors<BVSOA_t> interiors, InteriorBuildAid aid) {
		int idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= leafSize) return;

		int parent;
		if (idx == 0 || (idx != leafSize - 1 && aid.metrics[idx - 1] > aid.metrics[idx])) {
			// is left child of Interior idx
			parent = idx;
			leaves.parent[idx] = parent;
			interiors.lc[parent] = idx;
			interiors.rangeL[parent] = idx;
			//printf("interior %d has rangeL %d\n", parent, idx);
			leaves.segLen[idx] = 1;
			aid.leftLeafCount[idx] = 1;
			//atomicAdd(leaves.segLen + idx, 1);
			//atomicAdd(aid.leftLeafCount + idx, 1);
		}
		else {
			// is right child of Interior idx-1
			parent = idx - 1;
			leaves.parent[idx] = parent;
			interiors.rc[parent] = idx;
			interiors.rangeR[parent] = idx;
			//printf("interior %d has rangeR %d\n", parent, idx);
			//printf("buildInteriors: rangeR of %d = %d\n", parent, idx);
		}

		// choose parent in bottom-up fashion. O(n)
		int current, left, right;
		int _lc, _rc, lc, rc;
		bool isleaf_lc, isleaf_rc;
		while (atomicAdd(aid.visitCount + parent, 1) == 1) {
			current = parent;
			// merge bv
			_lc = interiors.lc[current], _rc = interiors.rc[current];
			isleaf_lc = isLeaf(_lc, &lc);
			isleaf_rc = isLeaf(_rc, &rc);
			if (!isleaf_lc && !isleaf_rc) {
				interiors.bv.setBV(current, interiors.bv, lc, interiors.bv, rc);
			}
			else if (!isleaf_lc && isleaf_rc) {
				interiors.bv.setBV(current, interiors.bv, lc, leaves.bv, rc);
			}
			else if (isleaf_lc && !isleaf_rc) {
				interiors.bv.setBV(current, leaves.bv, lc, interiors.bv, rc);
			}
			else {
				interiors.bv.setBV(current, leaves.bv, lc, leaves.bv, rc);
			}

			left = interiors.rangeL[current];
			right = interiors.rangeR[current];
			//if (right==0) printf("rangeR of %d in buildInteriors = 0\n", current);
			__threadfence();

			if (left == 0 && right == leafSize - 1) {
				//root
				break;
			}

			if (left == 0 || (right != leafSize - 1 && aid.metrics[left - 1] > aid.metrics[right])) {
				// is left child of Interior right
				parent = right;
				interiors.parent[current] = parent;
				interiors.lc[parent] = interiorIdx(current);
				interiors.rangeL[parent] = left;
				//printf("interior %d has rangeL %d\n", parent, left);

				//aid.leftLeafCount[parent] = old;
				leaves.segLen[left] ++;  // add LCL value of the leftmost leaf
				//atomicAdd(leaves.segLen + left, 1);
				aid.leftLeafCount[parent] = leaves.segLen[left]; // count the order of Interior parent
				__threadfence();
			}
			else {
				// is right child of Interior left-1
				parent = left - 1;
				interiors.parent[current] = parent;
				interiors.rc[parent] = interiorIdx(current);
				interiors.rangeR[parent] = right;
				__threadfence();
			}
		}
	}

	template<typename BVSOA_t>
	__global__ void calcInteriorNewIdx(int size, LBvhLeaves<BVSOA_t> leaves, LBvhInteriors<BVSOA_t> interiors, int* leftLeafCount, int* mapidx) {
		int idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= size) return;

		int binIdx = interiors.rangeL[idx];
		int idx_new = leaves.segOffset[binIdx] + leaves.segLen[binIdx] - leftLeafCount[idx];
		mapidx[idx] = idx_new;
	}

	template<typename BVSOA_t>
	__global__ void reorderInteriors(int interior_size, LBvhLeaves<BVSOA_t> leaves, LBvhInteriors<BVSOA_t> interiors, LBvhInteriors<BVSOA_t> interiors_sorted, int* mapidx) {
		int idx = blockIdx.x * blockDim.x + threadIdx.x;

		if (idx >= interior_size) return;
		int idx_new = mapidx[idx];

		interiors_sorted.rangeL[idx_new] = interiors.rangeL[idx];
		interiors_sorted.rangeR[idx_new] = interiors.rangeR[idx];
		auto bv = interiors.bv.getBV(idx);
		interiors_sorted.bv.setBV(idx_new, bv);
	}
}
}
}